package com.example.ai.embedding;

import com.example.ai.etl.MyKeywordEnricher;
import com.example.ai.etl.reader.MyPagePdfDocumentReader;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Component;
import org.springframework.web.bind.annotation.*;

import java.util.List;
import java.util.Map;

@RestController
@RequestMapping("/verctor")
public class ElasticSearchController {
    @Autowired
    VectorStore vectorStore;
    @Autowired
    MyPagePdfDocumentReader myPagePdfDocumentReader;

    @Autowired
    MyKeywordEnricher keywordEnricher;
    @GetMapping("/init/esData")
    public ResponseEntity<?> initEsData() {
        List<Document> docs = myPagePdfDocumentReader.getDocs("data/中华人民共和国劳动法.pdf");
        List<Document> documents = keywordEnricher.enrichDocuments(docs,3);
        vectorStore.add(documents);
        return ResponseEntity.ok("success");

    }

    @PostMapping("/es/similarSearch")
    public ResponseEntity<?> similarSearch(@RequestBody  SimilarSearchReq req) {
        // Retrieve documents similar to a query
        List<Document> results = this.vectorStore.similaritySearch(SearchRequest.builder()
                .query(req.getContent())
                .topK(req.getTopK())
                .similarityThreshold(req.getThreshold())
                .filterExpression(req.getFilter())
                .build());
        return ResponseEntity.ok(results);
    }

}
